#include "PCFG.h"
#include <mpi.h>
#include <vector>
#include <string>
#include <chrono>
#include <iostream>
#include "utils.h"
#include "md5.h"
using namespace std;
using namespace chrono;




int main(int argc, char** argv) {
    // 初始化MPI环境
    MPI_Init(&argc, &argv);
    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    PriorityQueue q;
    double time_hash = 0.0;    // 哈希计算总时间
    double time_guess = 0.0;   // 猜测生成总时间
    double time_train = 0.0;   // 模型训练总时间
    int total_guesses = 0;     // 总猜测数

    auto start_total = system_clock::now();

    // 主进程负责训练模型
    if (rank == 0) {
        auto start_train = system_clock::now();
        q.m.train("/guessdata/Rockyou-singleLined-full.txt");
        q.m.order();
        auto end_train = system_clock::now();
        auto duration_train = duration_cast<microseconds>(end_train - start_train);
        time_train = duration_train.count() * 1e-6;
        cout << "Training time: " << time_train << " seconds" << endl;
    }

    // 广播模型到所有进程
    q.m.broadcast();

    // 初始化优先队列
    q.init();

    if (rank == 0) {
        // 主进程逻辑
        cout << "Starting guessing with " << size << " processes" << endl;
        const int MAX_GUESSES = 10000000; // 最大生成猜测数
        const int PRINT_INTERVAL = 1000000; // 进度打印间隔

        auto start_guess = system_clock::now();
        int curr_guesses = 0;

        while (!q.priority.empty() && total_guesses < MAX_GUESSES) {
            // 从队列中取出最高概率的PT
            q.PopNext();

            // 记录当前生成的猜测数
            curr_guesses = q.guesses.size();
            total_guesses += curr_guesses;

            // 打印进度
            if (total_guesses % PRINT_INTERVAL == 0) {
                cout << "Generated " << total_guesses << " guesses" << endl;
            }

            // 模拟哈希计算时间（实际应用中应替换为真实哈希计算）
            auto start_hash = system_clock::now();
            bit32 state[4];
            for (const string& guess : q.guesses) {
                // 模拟MD5计算
                MD5Hash(guess, state);
            }
            auto end_hash = system_clock::now();
            time_hash += duration_cast<microseconds>(end_hash - start_hash).count() * 1e-6;

            // 清空猜测缓冲区
            q.guesses.clear();

            // 向工作进程分发任务
            for (int worker = 1; worker < size; ++worker) {
                if (!q.priority.empty()) {
                    // 发送任务标识
                    int has_task = 1;
                    MPI_Send(&has_task, 1, MPI_INT, worker, 0, MPI_COMM_WORLD);

                    // 发送当前处理的PT信息
                    PT current_pt = q.priority.front();
                    send_string(to_string(current_pt.pivot), worker, 1, MPI_COMM_WORLD);
                }
                else {
                    // 没有更多任务，通知工作进程退出
                    int has_task = 0;
                    MPI_Send(&has_task, 1, MPI_INT, worker, 0, MPI_COMM_WORLD);
                }
            }

            // 接收工作进程的结果并统计哈希时间
            for (int worker = 1; worker < size; ++worker) {
                int result_count;
                MPI_Recv(&result_count, 1, MPI_INT, worker, 2, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

                // 接收结果并模拟哈希计算
                start_hash = system_clock::now();
                bit32 state[4];
                for (int i = 0; i < result_count; ++i) {
                    string guess = recv_string(worker, 3, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

                    MD5Hash(guess, state);
                }
                end_hash = system_clock::now();
                time_hash += duration_cast<microseconds>(end_hash - start_hash).count() * 1e-6;
            }
        }

        auto end_guess = system_clock::now();
        time_guess = duration_cast<microseconds>(end_guess - start_guess).count() * 1e-6;

        // 计算总时间
        auto end_total = system_clock::now();
        double duration_total = duration_cast<microseconds>(end_total - start_total).count() * 1e-6;

        // 输出详细时间统计
        cout << "===== Performance Statistics =====" << endl;
        cout << "Total time: " << duration_total << " seconds" << endl;
        cout << "Train time: " << time_train << " seconds" << endl;
        cout << "Guess time: " << time_guess - time_hash << " seconds" << endl;
        cout << "Hash time: " << time_hash << " seconds" << endl;
        cout << "Total guesses generated: " << total_guesses << endl;
        cout << "Guesses per second: " << total_guesses / (time_guess - time_hash) << endl;
    }
    else {
        // 工作进程逻辑
        while (true) {
            // 接收任务标识
            int has_task;
            MPI_Recv(&has_task, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

            if (!has_task) break; // 没有更多任务，退出循环

            // 接收PT信息
            string pivot_str = recv_string(0, 1, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            int pivot = stoi(pivot_str);

            // 基于接收到的信息生成猜测
            vector<string> local_guesses;
            // 实际应用中，这里应该根据PT信息生成猜测
            for (int i = 0; i < 1000; ++i) {
                local_guesses.push_back("guess_" + to_string(rank) + "_" + to_string(i));
            }

            // 发送结果回主进程
            int result_count = local_guesses.size();
            MPI_Send(&result_count, 1, MPI_INT, 0, 2, MPI_COMM_WORLD);

            for (const string& guess : local_guesses) {
                send_string(guess, 0, 3, MPI_COMM_WORLD);
            }
        }
    }

    // 清理MPI环境
    MPI_Finalize();
    return 0;
}
